#!/usr/bin/env julia

"""
Radiance Observation Parity Check Script

Validates Julia radiance BUFR reading against Fortran GSI reference.
Part of E.8 regression harness implementation.

Usage:
    julia scripts/julia/diagnostics/radiance_parity_check.jl \\
        results/radiance_validation \\
        /path/to/fortran/reference

Exits with code 0 if all parity checks pass, non-zero otherwise.
"""

using JSON3
using Printf
using Statistics

const PARITY_THRESHOLDS = Dict(
    :obs_count_tolerance => 0.05,        # 5% drift in observation counts
    :channel_count_tolerance => 0.0,     # Exact channel count match required
    :geolocation_tolerance => 0.1,       # 0.1 degree geolocation tolerance
    :brightness_temp_range_tolerance => 0.05  # 5% range tolerance
)

struct ParityResult
    instrument::String
    metric::String
    julia_value::Any
    fortran_value::Any
    tolerance::Float64
    passed::Bool
    deviation::Float64
end

function load_julia_radiance_stats(results_dir::String)
    """Load Julia radiance statistics from JSON files."""
    stats = Dict{String,Any}()

    # Try both naming patterns: *_stats.json and *_stats_test.json
    for (instrument, filenames) in [
        ("ATMS", ["atms_stats.json", "atms_stats_test.json"]),
        ("MHS", ["mhs_stats.json", "mhs_stats_test.json"]),
        ("CrIS", ["cris_stats.json", "cris_stats_test.json"])
    ]
        filepath = ""
        for filename in filenames
            candidate = joinpath(results_dir, filename)
            if isfile(candidate)
                filepath = candidate
                break
            end
        end

        if isempty(filepath)
            @warn "Julia stats file not found for instrument" instrument
            continue
        end

        try
            # Read JSON - handle CrIS NaN values by allowing invalid JSON
            json_str = read(filepath, String)
            # Replace NaN with null for valid JSON
            json_str = replace(json_str, r":NaN([,}])" => s":null\1")

            data = JSON3.read(json_str)
            # Convert to Dict with String keys to avoid type issues
            stats[instrument] = Dict{String,Any}(
                string(k) => v for (k, v) in pairs(data)
            )
            @info "Loaded Julia stats" instrument n_obs=get(stats[instrument], "n_observations", 0) n_channels=get(stats[instrument], "n_channels", 0)
        catch err
            @error "Failed to load Julia stats" instrument error=err
        end
    end

    return stats
end

function parse_fortran_debufr(debufr_file::String)
    """
    Parse Fortran debufr output to extract observation counts.

    Note: This is a simplified parser. Full implementation would parse
    brightness temperatures and channels from the debufr text format.
    """
    if !isfile(debufr_file)
        @warn "Fortran debufr file not found" debufr_file
        return nothing
    end

    # Count non-empty lines as a proxy for observation count
    # Real implementation would parse the actual debufr format
    n_lines = 0
    open(debufr_file) do f
        for line in eachline(f)
            if !isempty(strip(line))
                n_lines += 1
            end
        end
    end

    return Dict(
        "n_observations" => n_lines,
        "source_file" => debufr_file
    )
end

function load_fortran_reference(fortran_dir::String)
    """Load Fortran GSI radiance reference data."""
    reference = Dict{String,Any}()

    for (instrument, filename) in [
        ("ATMS", "gdas.t00z.atms.tm00.bufr_d.debufr.out"),
        ("MHS", "gdas.t00z.1bmhs.tm00.bufr_d.debufr.out"),
        ("CrIS", "gdas.t00z.crisf4.tm00.bufr_d.debufr.out")
    ]
        filepath = joinpath(fortran_dir, filename)
        data = parse_fortran_debufr(filepath)
        if data !== nothing
            reference[instrument] = data
            @info "Loaded Fortran reference" instrument n_obs=data["n_observations"]
        end
    end

    return reference
end

function check_observation_count_parity(instrument::String, julia_stats::Dict,
                                        fortran_ref::Dict, sample_limit::Int)
    """
    Check observation count parity between Julia and Fortran.

    Note: Julia samples up to `sample_limit` observations, while Fortran
    debufr output contains all observations. We verify Julia read the
    expected sample size correctly.
    """
    julia_count = julia_stats["n_observations"]
    expected_count = min(sample_limit, fortran_ref["n_observations"])

    deviation = abs(julia_count - expected_count) / max(expected_count, 1)
    tolerance = PARITY_THRESHOLDS[:obs_count_tolerance]
    passed = deviation <= tolerance

    return ParityResult(
        instrument,
        "observation_count",
        julia_count,
        expected_count,
        tolerance,
        passed,
        deviation
    )
end

function check_channel_count_parity(instrument::String, julia_stats::Dict,
                                     expected_channels::Int)
    """Verify Julia decoded the expected number of channels."""
    julia_channels = julia_stats["n_channels"]

    deviation = abs(julia_channels - expected_channels) / max(expected_channels, 1)
    tolerance = PARITY_THRESHOLDS[:channel_count_tolerance]
    passed = deviation <= tolerance

    return ParityResult(
        instrument,
        "channel_count",
        julia_channels,
        expected_channels,
        tolerance,
        passed,
        deviation
    )
end

function check_geolocation_parity(instrument::String, julia_stats::Dict)
    """
    Verify geolocation coordinates are within valid ranges.

    Full parity check would compare with Fortran, but we verify
    physical validity here as a sanity check.
    """
    geom = julia_stats["geometry"]

    lat_valid = -90 <= geom["latitude_min"] <= 90 && -90 <= geom["latitude_max"] <= 90
    lon_valid = -180 <= geom["longitude_min"] <= 180 && -180 <= geom["longitude_max"] <= 180
    zenith_valid = 0 <= geom["zenith_min"] <= 90 && 0 <= geom["zenith_max"] <= 90

    passed = lat_valid && lon_valid && zenith_valid

    return ParityResult(
        instrument,
        "geolocation_validity",
        "lat: [$(geom["latitude_min"]), $(geom["latitude_max"])], " *
        "lon: [$(geom["longitude_min"]), $(geom["longitude_max"])], " *
        "zenith: [$(geom["zenith_min"]), $(geom["zenith_max"])]",
        "lat: [-90, 90], lon: [-180, 180], zenith: [0, 90]",
        0.0,
        passed,
        0.0
    )
end

function check_brightness_temp_ranges(instrument::String, julia_stats::Dict)
    """
    Verify brightness temperatures are within physically reasonable ranges.

    Expected ranges (instrument-dependent):
    - Microwave (ATMS, MHS): 150-300 K
    - Infrared (CrIS): 180-320 K
    """
    expected_range = instrument == "CrIS" ? (180.0, 320.0) : (150.0, 300.0)

    # Check all channels
    all_valid = true
    invalid_channels = String[]

    if haskey(julia_stats, "channel_statistics")
        for (ch, stats) in julia_stats["channel_statistics"]
            mean_tb = stats["mean"]
            if mean_tb === nothing || (mean_tb isa Number && isnan(mean_tb))
                continue  # Skip channels with no data
            end

            if mean_tb < expected_range[1] || mean_tb > expected_range[2]
                all_valid = false
                push!(invalid_channels, "ch$ch: $(mean_tb)K")
            end
        end
    end

    return ParityResult(
        instrument,
        "brightness_temp_range",
        all_valid ? "all valid" : "invalid: $(join(invalid_channels, ", "))",
        "$(expected_range[1])-$(expected_range[2]) K",
        0.0,
        all_valid,
        all_valid ? 0.0 : 1.0
    )
end

function run_parity_checks(julia_dir::String, fortran_dir::String; sample_limit::Int=10000)
    """Run all parity checks and return results."""

    println("="^80)
    println("Radiance Observation Parity Check")
    println("="^80)
    println("Julia results:   $julia_dir")
    println("Fortran reference: $fortran_dir")
    println("Sample limit:    $sample_limit")
    println()

    julia_stats = load_julia_radiance_stats(julia_dir)
    fortran_ref = load_fortran_reference(fortran_dir)

    results = ParityResult[]

    # Expected channel counts from instrument specifications
    expected_channels = Dict(
        "ATMS" => 22,
        "MHS" => 5,
        "CrIS" => 555  # Can vary: 399, 555, 1305 depending on processing
    )

    for instrument in ["ATMS", "MHS", "CrIS"]
        if !haskey(julia_stats, instrument)
            @warn "Julia stats missing for instrument" instrument
            continue
        end

        # Check 1: Observation count (if Fortran reference available)
        if haskey(fortran_ref, instrument)
            push!(results, check_observation_count_parity(
                instrument, julia_stats[instrument],
                fortran_ref[instrument], sample_limit))
        end

        # Check 2: Channel count
        push!(results, check_channel_count_parity(
            instrument, julia_stats[instrument],
            expected_channels[instrument]))

        # Check 3: Geolocation validity
        push!(results, check_geolocation_parity(
            instrument, julia_stats[instrument]))

        # Check 4: Brightness temperature ranges
        push!(results, check_brightness_temp_ranges(
            instrument, julia_stats[instrument]))
    end

    return results
end

function print_parity_report(results::Vector{ParityResult})
    """Print formatted parity check report."""

    println("\nParity Check Results")
    println("="^80)
    println()

    passed_count = 0
    failed_count = 0

    for result in results
        status_symbol = result.passed ? "✓" : "✗"
        status_color = result.passed ? "\e[32m" : "\e[31m"
        reset_color = "\e[0m"

        println("$(status_color)$status_symbol$(reset_color) $(result.instrument) - $(result.metric)")
        println("    Julia:     $(result.julia_value)")
        println("    Expected:  $(result.fortran_value)")
        if result.tolerance > 0
            println("    Deviation: $(round(result.deviation * 100, digits=2))% (threshold: $(round(result.tolerance * 100, digits=2))%)")
        end
        println()

        if result.passed
            passed_count += 1
        else
            failed_count += 1
        end
    end

    println("="^80)
    println("Summary: $passed_count passed, $failed_count failed")
    println("="^80)

    return failed_count == 0
end

function main()
    if length(ARGS) < 1
        println("Usage: julia radiance_parity_check.jl <julia_results_dir> [fortran_ref_dir] [sample_limit]")
        println()
        println("Example:")
        println("  julia scripts/julia/diagnostics/radiance_parity_check.jl \\")
        println("        results/radiance_validation \\")
        println("        /home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/satellite_gdas \\")
        println("        10000")
        exit(1)
    end

    julia_dir = ARGS[1]
    fortran_dir = length(ARGS) >= 2 ? ARGS[2] : ""
    sample_limit = length(ARGS) >= 3 ? parse(Int, ARGS[3]) : 10000

    if !isdir(julia_dir)
        @error "Julia results directory not found" julia_dir
        exit(1)
    end

    if isempty(fortran_dir) || !isdir(fortran_dir)
        @warn "Fortran reference directory not provided or not found - skipping obs count checks" fortran_dir
        fortran_dir = julia_dir  # Use dummy path
    end

    results = run_parity_checks(julia_dir, fortran_dir; sample_limit=sample_limit)
    all_passed = print_parity_report(results)

    if all_passed
        println("\n✅ All parity checks PASSED")
        exit(0)
    else
        println("\n❌ Some parity checks FAILED")
        exit(1)
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end
